#pragma once
#include <assert.h>
#include <iostream>

namespace my_avl
{
	// AVL树我们只掌控插入,通过旋转保证它是一个AVL树
	// 左孩子 平衡因子-1
	// 右孩子 平衡因子+1
	template <typename K, typename V>
	struct AVL_Node
	{
		AVL_Node(K key, V val)
			: _left(nullptr),
			_right(nullptr),
			_parent(nullptr),
			_bf(0),
			_key(key),
			_val(val)
		{
		}
		AVL_Node<K, V>* _left;
		AVL_Node<K, V>* _right;
		AVL_Node<K, V>* _parent;
		K _key;
		V _val;
		int _bf;
	};
	template <typename K, typename V>
	class AVL
	{
	private:
		typedef AVL_Node<K, V> Node;
		typedef AVL<K, V> Self;
		Node* get_node(K key, V val)
		{
			return new Node(key, val);
		}

		void LRRotate(Node* cur)
		{
			Node* curR = cur->_right;
			int bf = curR->_bf; // 经过旋转后 curR的bf会更改需要提前记录！
			LRotate(curR);
			RRotate(curR);
			if (bf == -1)
			{
				cur->_bf = 0;
				curR->_bf = 0;
				cur->_parent->_bf = 1;
			}
			else if (bf == 1)
			{
				cur->_bf = -1;
				curR->_bf = 0;
				cur->_parent->_bf = 0;
			}
			else
			{
				cur->_bf = 0;
				curR->_bf = 0;
				cur->_parent->_bf = 0;
			}
		}
		void RLRotate(Node* cur)
		{
			Node* curL = cur->_left;
			int bf = curL->_bf; // 经过旋转后 curR的bf会更改需要提前记录！
			RRotate(curL);
			LRotate(curL);
			if (bf == -1)
			{
				cur->_bf = 1;
				curL->_bf = 0;
				cur->_parent = 0;
			}
			else if (bf == 1)
			{
				cur->_bf = 0;
				curL->_bf = -1;
				cur->_parent = 0;
			}
			else
			{
				cur->_bf = 0;
				curL->_bf = 0;
				cur->_parent = 0;
			}
		}
		// 旋转的时候需要维护要父子之间的联系
		// 孩子改动，孩子的父亲也需要跟着改动
		void LRotate(Node* cur)
		{
			Node* parent = cur->_parent;
			Node* curL = cur->_left;
			Node* pparent = parent->_parent;

			parent->_right = curL;
			// curL 可能不存在
			if (curL)
				curL->_parent = parent;

			cur->_left = parent;
			parent->_parent = cur;

			// 如果父亲节点是根节点，则祖父节点并不存在
			if (parent == _root)
			{
				_root = cur;
				_root->_parent = nullptr;
			}
			else
			{
				if (pparent->_left == parent)
					pparent->_left = cur;
				else
					pparent->_right = cur;
				cur->_parent = pparent;
			}

			cur->_bf = parent->_bf = 0;
		}

		// 旋转的时候需要维护要父子之间的联系
		// 孩子改动，孩子的父亲也需要跟着改动
		void RRotate(Node* cur)
		{
			Node* parent = cur->_parent;
			Node* curR = cur->_right;
			Node* pparent = parent->_parent;
			// 父亲一定是存在的，进入左旋的条件就是 父亲的平衡因子为-2
			// 祖父不一定存在，因为父亲可能是根节点

			// curR 可能不存在
			parent->_left = curR;
			if (curR)
				curR->_parent = parent;

			cur->_right = parent;
			parent->_parent = cur;

			// 如果父亲是根节点
			if (parent == _root)
			{
				_root = cur;
				_root->_parent = nullptr;
			}
			else
			{
				if (pparent->_left == parent)
					pparent->_left = cur;
				else
					pparent->_right = cur;
				cur->_parent = pparent;

			}

			cur->_bf = parent->_bf = 0;
		}

		int _GetHeight(Node* root)
		{
			if (root == nullptr)
				return 0;
			int leftH = _GetHeight(root->_left);
			int rightH = _GetHeight(root->_right);
			return std::max(leftH, rightH) + 1;
		}

		// AVL是一颗二叉搜索树
		// 高度差不超过1，平衡因子不能出错
		bool _IsAVL(Node* root, Node*& pre)
		{
			if (root == nullptr)
				return true;
			bool left = _IsAVL(root->_left, pre);

			int leftH = _GetHeight(root->_left);
			int rightH = _GetHeight(root->_right);
			if (!(rightH - leftH == 1 && root->_bf == 1 ||
				rightH - leftH == -1 && root->_bf == -1 ||
				rightH - leftH == 0 && root->_bf == 0))
			{
				std::cout << "高度差 or 平衡因子错误" << std::endl;
				std::cout << "leftH = " << leftH << " " << "rightH = " << rightH << std::endl;
				std::cout << "bf = " << root->_bf << std::endl;
				return false;
			}

			if (pre == nullptr)
				pre = root;
			else
			{
				if (root->_key <= pre->_key)
				{
					std::cout << "不是搜索二叉树 root->_key <= pre->_key" << std::endl;
					return false;
				}
				pre = root;
			}

			bool right = _IsAVL(root->_right, pre);
			return left && right;
		}
        
		void _PreOrder(Node* root)
		{
			if (root == nullptr)
			{
				std::cout << "# ";
				return;
			}
			std::cout << root->_key << " ";
			_PreOrder(root->_left);
			_PreOrder(root->_right);
		}
	public:
		void PreOrder()
		{
			_PreOrder(_root);
		}
		bool IsAVL()
		{
			Node* pre = nullptr;
			return _IsAVL(_root, pre);
		}
		bool Insert(K key, V val)
		{
			if (_root == nullptr)
			{
				_root = get_node(key, val);
				return true;
			}
			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				if (key > cur->_key)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (key < cur->_key)
				{
					parent = cur;
					cur = cur->_left;
				}
				else
					return false;
			}
			cur = get_node(key, val);
			cur->_parent = parent;
			if (key > parent->_key)
				parent->_right = cur;
			else
				parent->_left = cur;
			// 检测平衡因子
			while (parent)
			{

				if (parent->_left == cur)
					parent->_bf--;
				else
					parent->_bf++;

				if (cur->_bf == -1 && parent->_bf == -2)
				{
					RRotate(cur);
					break;
				}
				else if (cur->_bf == 1 && parent->_bf == 2)
				{
					LRotate(cur);
					break;
				}
				else if (cur->_bf == -1 && parent->_bf == 2)
				{
					RLRotate(cur);
					break;
				}
				else if (cur->_bf == 1 && parent->_bf == -2)
				{
					LRRotate(cur);
					break;
				}
				else if (parent->_bf > 2 && parent->_bf < -2)
					assert(0);
				cur = parent;
				parent = parent->_parent;
			}
			return true;
		}

	private:
		Node* _root = nullptr;
	};
}